"""
将输入文本（单条查询或多条文档）转换为固定维度的向量表示，并对向量进行归一化处理，确保向量具有统一的尺度（便于后续相似度计算）。
texts = [
    "被告人张某因故意伤人罪被判处有期徒刑 3 年，罚金 2000 元。",
    "非法狩猎野生动物属于违法行为，情节严重者将面临刑事处罚。",
    "检察院指控被告人李某犯诈骗罪，涉案金额共计10万元。"
]
=>
[
    [0.01234567, -0.04567890, 0.07890123, ...],  # 第一条文本的向量（1024 维）
    [0.05678901, -0.02345678, 0.09876543, ...],  # 第二条文本的向量（1024 维）
    [0.03456789, -0.06789012, 0.01234567, ...]   # 第三条文本的向量（1024 维）
]
"""
from langchain.embeddings.base import Embeddings
from typing import List
import requests
import json
from transformers import AutoModel, AutoTokenizer
import torch


# 标准化向量
def normal(vector):
	ss = sum([s ** 2 for s in vector]) ** 0.5  # 计算向量的 L2 范数（模长）
	return [round(s / ss, 8) for s in vector]  # 每个元素除以模长，保留 8 位小数


class CustomEmbeddings(Embeddings):
	def __init__(self):
		pass
	
	def embed_documents(self, texts: List[str]) -> List[List[float]]:
		"""
		功能：批量处理多条文档文本，返回一个向量列表（每条文档对应一个向量）。
		"""
		results = [self.embed_query(text) for text in texts]
		return results
	
	def embed_query(self, text: str) -> List[float]:
		"""
		功能：处理单条查询文本，返回其归一化后的向量表示。
		"""
		# 文本分词：转换为模型可接受的token ID，自动 padding 和截断，返回 PyTorch 张量并转移到 GPU
		token = tokenizer([text], padding=True, truncation=True, return_tensors='pt').to(device)
		# 模型推理：获取 [CLS] 位置的输出向量（BERT 模型的第 2 个返回值，即池化后的句子向量）
		# model_weight(**token) == model_weight(input_ids=token["input_ids"], attention_mask=token["attention_mask"])
		vector = model(**token)[1].tolist()[0]
		# 向量归一化
		vector = normal(vector)
		return vector


count = 0
device = torch.device("cuda:0")
# 标准 bert 模型，最为向量模型
# model_name = "E:\\code\\bge\\bge_recall"
model_name = "BAAI/bge-large-zh"  # BAAI/bge-large-zh 是一个由北京智源人工智能研究院（BAAI）开发的中文文本嵌入模型。
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name).to(device)
# print("模型加载成功")
# print(model_weight)
